from omegaconf import DictConfig
from lightning import Trainer
from lightning.pytorch.loggers import WandbLogger

from ecgcmr.imaging.img_dataset.DownstreamImagingLightning import DownstreamImageDataModule
from ecgcmr.imaging.img_models.ImageViTMAEEval import ImageViTEval
from ecgcmr.imaging.img_models.ImageSupervised_ResNet3D import ImageEncoderSupervised_ResNet3D


def train_imaging_supervised(
    cfg: DictConfig,
    wandb_logger: WandbLogger,
    save_dir: str,
    devices: int = 1
    ):
    datamodule = DownstreamImageDataModule(cfg=cfg, mask_labels=False, supervised=True)

    if cfg.models.backbone == 'vit':
        model = ImageViTEval(cfg=cfg, save_dir=save_dir)
    elif cfg.models.backbone == 'resnet':
        model = ImageEncoderSupervised_ResNet3D(cfg=cfg, save_dir=save_dir)
        
    wandb_logger.watch(model, log_graph=False)

    strategy = "ddp" if devices > 1 else "auto"

    trainer = Trainer(accelerator="gpu",
                      devices=devices,
                      precision="bf16-mixed",
                      logger=wandb_logger,
                      strategy=strategy,
                      max_epochs=cfg.max_epochs,
                      log_every_n_steps=cfg.log_every_n_steps,
                      check_val_every_n_epoch=cfg.check_val_every_n_epoch,
                      default_root_dir=save_dir,
                      num_sanity_val_steps=0,
                      profiler="simple"
                      )
    
    trainer.fit(model=model, datamodule=datamodule)